#include "Fragment_TransportLayer.h"
#include "Fragment_LinkLayer.h"
#include "string.h"


static Fragment_ReassemblyContext_t xReassemblyContext;
static TransferFunc_t pfnTransmit = NULL;
static TransferFunc_t pfnReceivedCmp = NULL;


static uint16_t Fragment_checksum(const uint8_t* data, uint16_t length) 
{
    if (data == NULL || length < 2) return 0;
    uint16_t sum_value = 0;
    for (uint16_t i = 0; i < length; i++) {
        sum_value += data[i];
    }
    return sum_value;
}

static uint16_t Fragment_GetPacketId(void) 
{
    static uint8_t frame_id = 0;
    return frame_id++;
}

bool Fragment_Packet_Split_Init(Fragment_SplitContext_t* ctx, const uint8_t* data, uint16_t length)
{
    if (!ctx || !data || length == 0 || length > MAX_TOTAL_LENGTH) return false;
    memset(ctx, 0, sizeof(Fragment_SplitContext_t));
    ctx->data = data;
    ctx->checksum = Fragment_checksum(data, length); // 计算校验和
    ctx->total_length = length + 2;  // exclude checksum
    ctx->packet_id = Fragment_GetPacketId(); // 初始包ID
    ctx->total_packets = (ctx->total_length + MAX_PKT_PAYLOAD_LEN - 1) / MAX_PKT_PAYLOAD_LEN;
    ctx->current_index = 0; // 初始化当前索引
    return true;
}


bool Fragment_Packet_Split(Fragment_SplitContext_t* ctx, Fragment_Packet_t* out_packet)
{
    if (!ctx || !out_packet || !ctx->data) return false;
    if (ctx->current_index >= ctx->total_packets) return false;

    uint16_t offset = ctx->current_index * MAX_PKT_PAYLOAD_LEN;
    uint16_t remain_total_length = ctx->total_length - offset; // include checksum
    uint16_t remain_data_length = 0;                           // exclude checksum
    uint8_t chunk_len = 0;                                     // copy data length (exclude checksum)
    if (remain_total_length < 2) {
        remain_data_length = 0;
    } else {
        remain_data_length = remain_total_length - 2;
    }
    chunk_len = (remain_data_length > MAX_PKT_PAYLOAD_LEN) ? MAX_PKT_PAYLOAD_LEN : remain_data_length;
    
    out_packet->packet_id = ctx->packet_id;
    out_packet->total_packets = ctx->total_packets;
    out_packet->current_index = ctx->current_index;
    out_packet->payload_length = chunk_len;

    if (chunk_len) {
        memcpy(out_packet->payload, &ctx->data[offset], chunk_len);
        out_packet->payload_length = chunk_len;
    }
    

    remain_total_length -= chunk_len;
    while ((out_packet->payload_length < MAX_PKT_PAYLOAD_LEN) && remain_total_length)
    {
        if (remain_total_length > 2) break;

        if (remain_total_length == 1) {
            out_packet->payload[out_packet->payload_length] = (ctx->checksum >> 8) & 0xFF; // Write MSB first
        } else {
            out_packet->payload[out_packet->payload_length] = ctx->checksum & 0xFF; // Write LSB
        }
        
        out_packet->payload_length++;
        remain_total_length--;
    }
    
    ctx->current_index++;
    return true;
}


bool Fragment_Packet_Reassemble_Init(Fragment_ReassemblyContext_t* ctx)
{
    if (!ctx) return false;
    memset(ctx, 0, sizeof(Fragment_ReassemblyContext_t));
    return true;
}

extern void hex_dump(const char* label, const void* data, size_t len);

Fragment_ReassemblyState_e Fragment_Packet_Reassemble(Fragment_ReassemblyContext_t* ctx, const Fragment_Packet_t* pkt)
{
    if (!ctx || !pkt || (pkt->current_index >= MAX_TOTAL_PACKETS)) return ReassemblyResult_NuLLPtr;

    // 第一次包，初始化上下文
    if ((ctx->received_count == 0) || (pkt->current_index == 0))
    {
        Fragment_Packet_Reassemble_Init(ctx);
        ctx->total_packets = pkt->total_packets;
        ctx->packet_id = pkt->packet_id;
    }

    if ((ctx->packet_id != pkt->packet_id) && (pkt->current_index != 0) ) return ReassemblyResult_PacketIdMismatch;

    // 检查是否已接收
    if (BITMAP_GET(ctx->received_bitmap, pkt->current_index)) return ReassemblyResult_RepeatFrame;

    // 拷贝数据
    uint16_t offset = pkt->current_index * MAX_PKT_PAYLOAD_LEN;
    if (offset + pkt->payload_length > MAX_TOTAL_LENGTH_CHECKSUM) return ReassemblyResult_Overflow;

    memcpy(&ctx->whole_payload[offset], pkt->payload, pkt->payload_length);

    // 标记已接收
    BITMAP_SET(ctx->received_bitmap, pkt->current_index);
    ctx->received_count++;


    // 是否接收完
    if (ctx->received_count == ctx->total_packets) {
        ctx->whole_length = (pkt->total_packets - 1) * MAX_PKT_PAYLOAD_LEN + pkt->payload_length;

        if (ctx->whole_length > MAX_TOTAL_LENGTH_CHECKSUM || ctx->whole_length < 2) return ReassemblyResult_Overflow;

        
        ctx->checksum = ctx->whole_payload[ctx->whole_length - 2] | \
                        (ctx->whole_payload[ctx->whole_length - 1] << 8);

        if (ctx->checksum != Fragment_checksum(&ctx->whole_payload[0], ctx->whole_length - 2)) return ReassemblyResult_ChecksumFail;
        ctx->whole_length -= 2; // exclude checksum
        ctx->flag_complete = true;
        ctx->received_count = 0; // 清空状态（可选）
    }

    return ReassemblyResult_Successful;
}


uint16_t Fragment_Packet_Serialize(const Fragment_Packet_t* in_Pkt, uint8_t* output, uint16_t buf_size)
{
    if (!in_Pkt || !output || buf_size < (in_Pkt->payload_length + 4)) return 0;

    uint16_t index = 0;
    output[index++] = in_Pkt->packet_id;
    output[index++] = in_Pkt->total_packets;
    output[index++] = in_Pkt->current_index;
    output[index++] = in_Pkt->payload_length;
    memcpy(&output[index], in_Pkt->payload, in_Pkt->payload_length);
    index += in_Pkt->payload_length;
    return index;
}

bool Fragment_Packet_Parse(const uint8_t* in_Buff, const uint16_t length, Fragment_Packet_t* out_Pkt)
{
    if (!in_Buff || !out_Pkt || length < 4) return 0;

    uint16_t index = 0;
    out_Pkt->packet_id = in_Buff[index++];
    out_Pkt->total_packets = in_Buff[index++];
    out_Pkt->current_index = in_Buff[index++];
    out_Pkt->payload_length = in_Buff[index++];
    if (out_Pkt->payload_length > MAX_PKT_PAYLOAD_LEN || (length - index) < out_Pkt->payload_length) return false;
    memcpy(out_Pkt->payload, &in_Buff[index], out_Pkt->payload_length);
    
    return true;
}




void Fragment_Transmitter(const uint8_t* in_Buff, const uint16_t length)
{
    Fragment_SplitContext_t xSplitContext;

    Fragment_Packet_t xPacket;
    bool result = false;

    result = Fragment_Packet_Split_Init(&xSplitContext, in_Buff, length);
    if (result == false) return;

    
    FrameBuff_t xFramePayloadBuff = {0};
    FrameBuff_t xFrameBuff = {0};

    for (uint16_t i = 0; i < xSplitContext.total_packets; i++) {
        result = Fragment_Packet_Split(&xSplitContext, &xPacket);


        xFramePayloadBuff.length = Fragment_Packet_Serialize(&xPacket, xFramePayloadBuff.arrData, MAX_FRAME_LEN);


        xFrameBuff.length = Fragment_Frame_Serialize(&xFramePayloadBuff, xFrameBuff.arrData, MAX_FRAME_LEN);

        if (pfnTransmit) {
            pfnTransmit(xFrameBuff.arrData, xFrameBuff.length);
        }
    }
}

void Fragment_Received(const uint8_t* in_Buff, const uint16_t length)
{
    if (!in_Buff || length == 0) return;
    Fragment_LinkError_e link_result = LINK_SUCCESS;
    FrameBuff_t xFramePayloadBuff = {0};
    Fragment_Packet_t xPacket = {0};

    link_result = Fragment_Frame_Parse(in_Buff, length, &xFramePayloadBuff);

    if (link_result != LINK_SUCCESS) return;
    bool packet_parse_result = false;
    packet_parse_result = Fragment_Packet_Parse(xFramePayloadBuff.arrData, xFramePayloadBuff.length, &xPacket);

    if (packet_parse_result == false) return;
    
    Fragment_ReassemblyState_e reassembly_result = Fragment_Packet_Reassemble(&xReassemblyContext, &xPacket);

    
    if (xReassemblyContext.flag_complete && pfnReceivedCmp) {
        xReassemblyContext.flag_complete = false;
        pfnReceivedCmp(xReassemblyContext.whole_payload, xReassemblyContext.whole_length);
    }

}


void Fragment_TransportInit(TransferFunc_t send_callback, TransferFunc_t recv_callback)
{
    Fragment_Packet_Reassemble_Init(&xReassemblyContext);

    pfnTransmit = send_callback;
    pfnReceivedCmp = recv_callback;
}

